import logging
from typing import Dict, List, Optional

from config.settings import (GRAPH_DB_PASSWORD, GRAPH_DB_URL, GRAPH_DB_USER,
                            USE_GRAPH_DB)

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class GraphDBManager:
    """
    管理代码关系的知识图谱数据库
    """
    def __init__(self):
        self.enabled = USE_GRAPH_DB
        self.graph = None
        
        if self.enabled:
            try:
                # 条件导入，仅当启用图数据库时才加载依赖
                from neo4j import GraphDatabase
                self.driver = GraphDatabase.driver(
                    GRAPH_DB_URL, 
                    auth=(GRAPH_DB_USER, GRAPH_DB_PASSWORD)
                )
                logger.info("成功连接到Neo4j图数据库")
            except Exception as e:
                self.enabled = False
                logger.error(f"连接Neo4j数据库失败: {str(e)}")
                logger.info("将使用内存存储分析结果")
        else:
            logger.info("图数据库功能已禁用，使用内存存储")
            
        # 内存存储备份
        self.file_nodes = {}
        self.function_nodes = {}
        self.call_relations = []

    def close(self):
        """关闭数据库连接"""
        if self.enabled and self.driver:
            self.driver.close()

    def store_analysis_results(self, analysis_results: Dict) -> bool:
        """
        将代码分析结果存储到知识图谱
        """
        try:
            # 首先存储到内存中作为备份
            self._store_in_memory(analysis_results)
            
            # 如果启用了图数据库，则同时存储到Neo4j
            if self.enabled:
                self._store_in_neo4j(analysis_results)
                
            return True
        except Exception as e:
            logger.error(f"存储分析结果时出错: {str(e)}")
            return False

    def _store_in_memory(self, analysis_results: Dict) -> None:
        """将分析结果存储到内存"""
        call_graph = analysis_results.get("call_graph", {})
        file_functions = analysis_results.get("file_functions", {})
        
        # 存储文件节点
        for file_path in call_graph.keys():
            self.file_nodes[file_path] = {
                "path": file_path
            }
        
        # 存储函数节点及调用关系
        for file_path, functions in call_graph.items():
            for func_name, calls in functions.items():
                func_id = f"{file_path}:{func_name}"
                
                # 获取函数详情
                func_details = {}
                for func_info in file_functions.get(file_path, []):
                    if func_info.get("name") == func_name:
                        func_details = func_info
                        break
                
                self.function_nodes[func_id] = {
                    "name": func_name,
                    "file_path": file_path,
                    "params": func_details.get("params", []),
                    "return_type": func_details.get("return_type", ""),
                    "docstring": func_details.get("docstring", "")
                }
                
                # 存储调用关系
                for call in calls:
                    if call["type"] == "local":
                        target_func_id = f"{call['file']}:{call['name']}"
                        self.call_relations.append({
                            "source": func_id,
                            "target": target_func_id,
                            "type": "CALLS"
                        })
                    elif call["type"] == "external":
                        # 外部调用使用模块名作为标识
                        target_id = f"{call['module']}:{call['method']}"
                        self.call_relations.append({
                            "source": func_id,
                            "target": target_id,
                            "type": "CALLS_EXTERNAL"
                        })
        
        logger.info(f"已在内存中存储 {len(self.file_nodes)} 个文件和 {len(self.function_nodes)} 个函数")

    def _store_in_neo4j(self, analysis_results: Dict) -> None:
        """将分析结果存储到Neo4j图数据库"""
        if not self.enabled:
            return
            
        with self.driver.session() as session:
            # 清空现有数据
            session.run("MATCH (n) DETACH DELETE n")
            
            # 批处理参数，提高写入性能
            file_params = []
            function_params = []
            relation_params = []
            
            for file_id, file_data in self.file_nodes.items():
                file_params.append({"id": file_id, "path": file_data["path"]})
            
            for func_id, func_data in self.function_nodes.items():
                function_params.append({
                    "id": func_id,
                    "name": func_data["name"],
                    "file_id": func_data["file_path"],
                    "params": ",".join(func_data["params"]),
                    "return_type": func_data["return_type"],
                    "docstring": func_data["docstring"]
                })
            
            for rel in self.call_relations:
                relation_params.append({
                    "source": rel["source"],
                    "target": rel["target"],
                    "type": rel["type"]
                })
            
            # 批量创建文件节点
            if file_params:
                session.run("""
                    UNWIND $files AS file
                    CREATE (f:File {id: file.id, path: file.path})
                """, {"files": file_params})
            
            # 批量创建函数节点并关联到文件
            if function_params:
                session.run("""
                    UNWIND $functions AS func
                    MATCH (f:File {id: func.file_id})
                    CREATE (fn:Function {
                        id: func.id,
                        name: func.name,
                        params: func.params,
                        return_type: func.return_type,
                        docstring: func.docstring
                    })
                    CREATE (fn)-[:DEFINED_IN]->(f)
                """, {"functions": function_params})
            
            # 批量创建调用关系
            if relation_params:
                session.run("""
                    UNWIND $relations AS rel
                    MATCH (source {id: rel.source})
                    MATCH (target {id: rel.target})
                    CALL apoc.create.relationship(source, rel.type, {}, target)
                    YIELD rel as created
                    RETURN count(created)
                """, {"relations": relation_params})
                
            logger.info(f"已在Neo4j中存储 {len(file_params)} 个文件和 {len(function_params)} 个函数")

    def get_function_context(self, file_path: str, function_name: str, depth: int = 2) -> Dict:
        """
        获取函数的上下文信息，包括调用关系
        """
        func_id = f"{file_path}:{function_name}"
        
        if func_id not in self.function_nodes:
            return {"error": "Function not found"}
        
        # 基本函数信息
        context = dict(self.function_nodes[func_id])
        context["defined_in"] = file_path
        
        # 获取此函数调用的函数
        context["calls"] = []
        for rel in self.call_relations:
            if rel["source"] == func_id:
                called_func = rel["target"]
                if called_func in self.function_nodes:
                    call_info = {
                        "name": self.function_nodes[called_func]["name"],
                        "file": self.function_nodes[called_func]["file_path"]
                    }
                    context["calls"].append(call_info)
                else:
                    # 可能是外部库调用
                    module, method = called_func.split(":", 1)
                    context["calls"].append({
                        "name": method,
                        "module": module,
                        "external": True
                    })
        
        # 获取调用此函数的函数
        context["called_by"] = []
        for rel in self.call_relations:
            if rel["target"] == func_id and rel["type"] == "CALLS":
                caller_func = rel["source"]
                if caller_func in self.function_nodes:
                    caller_info = {
                        "name": self.function_nodes[caller_func]["name"],
                        "file": self.function_nodes[caller_func]["file_path"]
                    }
                    context["called_by"].append(caller_info)
        
        return context
